#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
股票快照数据表模型 (HSnapshot)
"""

from sqlalchemy import Column, Integer, String, Float, Index
from typing import List, Optional
import pandas as pd

from .BaseModel import BaseModel, DatabaseManager, safe_convert_numeric, safe_convert_int, safe_get_field_value, standardize_date_format

class HSnapshot(BaseModel):
    """股票快照数据表模型"""
    __tablename__ = 'hsnapshot'
    
    # 列名映射（CSV列名 -> 数据库字段名）
    aliases = {
        'tradeDate': 'trade_date',
        'market': 'market',
        'securityId': 'security_id',
        'preClosePx': 'pre_close_px',
        'openPx': 'open_px',
        'highPx': 'high_px',
        'lowPx': 'low_px',
        'tradePx': 'trade_px',
        'numTrades': 'num_trades',
        'volumeTrade': 'volume_trade',
        'amountTrade': 'amount_trade',
        'preVolume': 'pre_volume',
        'multVol': 'mult_vol',
        'upperTimes': 'upper_times',
        'upperBidOrdQty': 'upper_bid_ord_qty',
        'upperBidOrdCnt': 'upper_bid_ord_cnt',
        'lowerOfferOrdQty': 'lower_offer_ord_qty',
        'lowerOfferOrdCnt': 'lower_offer_ord_cnt',
        'upperOfferOrdQty': 'upper_offer_ord_qty',
        'upperOfferOrdCnt': 'upper_offer_ord_cnt',
        'lowerBidOrdQty': 'lower_bid_ord_qty',
        'lowerBidOrdCnt': 'lower_bid_ord_cnt',
        'yesUpperOfferOrdQty': 'yes_upper_offer_ord_qty',
        'yesUpperOfferOrdCnt': 'yes_upper_offer_ord_cnt',
        'upperPx': 'upper_px',
        'lowerPx': 'lower_px',
        'securityName': 'security_name',
        'netamt': 'net_amt',
        'bignetamt': 'big_net_amt',
        'biascnt': 'bias_cnt',
        'auctionamt': 'auction_amt',
        'publicfloatShare': 'public_float_share'
    }
    
    trade_date = Column('trade_date', String(10), primary_key=True, key='trade_date')
    market = Column('market', Integer, primary_key=True, key='market')
    security_id = Column('security_id', String(20), primary_key=True, key='security_id')
    pre_close_px = Column('pre_close_px', Float, key='pre_close_px')
    open_px = Column('open_px', Float, key='open_px')
    high_px = Column('high_px', Float, key='high_px')
    low_px = Column('low_px', Float, key='low_px')
    trade_px = Column('trade_px', Float, key='trade_px')
    num_trades = Column('num_trades', Integer, key='num_trades')
    volume_trade = Column('volume_trade', Integer, key='volume_trade')
    amount_trade = Column('amount_trade', Float, key='amount_trade')
    pre_volume = Column('pre_volume', Integer, key='pre_volume')
    mult_vol = Column('mult_vol', Integer, key='mult_vol')
    upper_times = Column('upper_times', Integer, key='upper_times')
    upper_bid_ord_qty = Column('upper_bid_ord_qty', Integer, key='upper_bid_ord_qty')
    upper_bid_ord_cnt = Column('upper_bid_ord_cnt', Integer, key='upper_bid_ord_cnt')
    lower_offer_ord_qty = Column('lower_offer_ord_qty', Integer, key='lower_offer_ord_qty')
    lower_offer_ord_cnt = Column('lower_offer_ord_cnt', Integer, key='lower_offer_ord_cnt')
    upper_offer_ord_qty = Column('upper_offer_ord_qty', Integer, key='upper_offer_ord_qty')
    upper_offer_ord_cnt = Column('upper_offer_ord_cnt', Integer, key='upper_offer_ord_cnt')
    lower_bid_ord_qty = Column('lower_bid_ord_qty', Integer, key='lower_bid_ord_qty')
    lower_bid_ord_cnt = Column('lower_bid_ord_cnt', Integer, key='lower_bid_ord_cnt')
    yes_upper_offer_ord_qty = Column('yes_upper_offer_ord_qty', Integer, key='yes_upper_offer_ord_qty')
    yes_upper_offer_ord_cnt = Column('yes_upper_offer_ord_cnt', Integer, key='yes_upper_offer_ord_cnt')
    upper_px = Column('upper_px', Float, key='upper_px')
    lower_px = Column('lower_px', Float, key='lower_px')
    security_name = Column('security_name', String(100), key='security_name')
    net_amt = Column('net_amt', Float, key='net_amt')
    big_net_amt = Column('big_net_amt', Float, key='big_net_amt')
    bias_cnt = Column('bias_cnt', Integer, key='bias_cnt')
    auction_amt = Column('auction_amt', Float, key='auction_amt')
    
    # 索引
    __table_args__ = (
        Index('idx_hsnapshot_trade_date', 'trade_date'),
        Index('idx_hsnapshot_security_id', 'security_id'),
        Index('idx_hsnapshot_trade_date_security_id', 'trade_date', 'security_id'),
        Index('idx_hsnapshot_market', 'market'),
    )

    @classmethod
    def get_session(cls):
        """获取数据库会话"""
        return cls._db_manager.get_stock_data_session()
    
    @classmethod
    def _get_existing_record(cls, session, data):
        """获取现有记录"""
        return session.query(cls).filter(
            cls.trade_date == data['trade_date'],
            cls.market == data['market'],
            cls.security_id == data['security_id']
        ).first()
    
    @classmethod
    def _update_existing_record(cls, existing, data):
        """更新现有记录"""
        existing.pre_close_px = data.get('pre_close_px', existing.pre_close_px)
        existing.open_px = data.get('open_px', existing.open_px)
        existing.high_px = data.get('high_px', existing.high_px)
        existing.low_px = data.get('low_px', existing.low_px)
        existing.trade_px = data.get('trade_px', existing.trade_px)
        existing.num_trades = data.get('num_trades', existing.num_trades)
        existing.volume_trade = data.get('volume_trade', existing.volume_trade)
        existing.amount_trade = data.get('amount_trade', existing.amount_trade)
        existing.pre_volume = data.get('pre_volume', existing.pre_volume)
        existing.mult_vol = data.get('mult_vol', existing.mult_vol)
        existing.upper_times = data.get('upper_times', existing.upper_times)
        existing.upper_bid_ord_qty = data.get('upper_bid_ord_qty', existing.upper_bid_ord_qty)
        existing.upper_bid_ord_cnt = data.get('upper_bid_ord_cnt', existing.upper_bid_ord_cnt)
        existing.lower_offer_ord_qty = data.get('lower_offer_ord_qty', existing.lower_offer_ord_qty)
        existing.lower_offer_ord_cnt = data.get('lower_offer_ord_cnt', existing.lower_offer_ord_cnt)
        existing.upper_offer_ord_qty = data.get('upper_offer_ord_qty', existing.upper_offer_ord_qty)
        existing.upper_offer_ord_cnt = data.get('upper_offer_ord_cnt', existing.upper_offer_ord_cnt)
        existing.lower_bid_ord_qty = data.get('lower_bid_ord_qty', existing.lower_bid_ord_qty)
        existing.lower_bid_ord_cnt = data.get('lower_bid_ord_cnt', existing.lower_bid_ord_cnt)
        existing.yes_upper_offer_ord_qty = data.get('yes_upper_offer_ord_qty', existing.yes_upper_offer_ord_qty)
        existing.yes_upper_offer_ord_cnt = data.get('yes_upper_offer_ord_cnt', existing.yes_upper_offer_ord_cnt)
        existing.upper_px = data.get('upper_px', existing.upper_px)
        existing.lower_px = data.get('lower_px', existing.lower_px)
        existing.security_name = data.get('security_name', existing.security_name)
        existing.net_amt = data.get('net_amt', existing.net_amt)
        existing.big_net_amt = data.get('big_net_amt', existing.big_net_amt)
        existing.bias_cnt = data.get('bias_cnt', existing.bias_cnt)
        existing.auction_amt = data.get('auction_amt', existing.auction_amt)
    
    @classmethod
    def _to_dict(cls, result):
        """将结果转换为字典"""
        return {
            'trade_date': result.trade_date,
            'market': result.market,
            'security_id': result.security_id,
            'pre_close_px': result.pre_close_px,
            'open_px': result.open_px,
            'high_px': result.high_px,
            'low_px': result.low_px,
            'trade_px': result.trade_px,
            'num_trades': result.num_trades,
            'volume_trade': result.volume_trade,
            'amount_trade': result.amount_trade,
            'pre_volume': result.pre_volume,
            'mult_vol': result.mult_vol,
            'upper_times': result.upper_times,
            'upper_bid_ord_qty': result.upper_bid_ord_qty,
            'upper_bid_ord_cnt': result.upper_bid_ord_cnt,
            'lower_offer_ord_qty': result.lower_offer_ord_qty,
            'lower_offer_ord_cnt': result.lower_offer_ord_cnt,
            'upper_offer_ord_qty': result.upper_offer_ord_qty,
            'upper_offer_ord_cnt': result.upper_offer_ord_cnt,
            'lower_bid_ord_qty': result.lower_bid_ord_qty,
            'lower_bid_ord_cnt': result.lower_bid_ord_cnt,
            'yes_upper_offer_ord_qty': result.yes_upper_offer_ord_qty,
            'yes_upper_offer_ord_cnt': result.yes_upper_offer_ord_cnt,
            'upper_px': result.upper_px,
            'lower_px': result.lower_px,
            'security_name': result.security_name,
            'net_amt': result.net_amt,
            'big_net_amt': result.big_net_amt,
            'bias_cnt': result.bias_cnt,
            'auction_amt': result.auction_amt
        }
    
    @classmethod
    def get_data(cls, trade_date: str, security_id: Optional[str] = None) -> List['HSnapshot']:
        """获取股票快照数据"""
        if not cls._db_manager:
            raise ValueError("数据库管理器未设置，请先调用 set_db_manager()")
        
        try:
            session = cls.get_session()
            
            query = session.query(cls).filter(cls.trade_date == trade_date)
            
            if security_id:
                query = query.filter(cls.security_id == security_id)
            
            results = query.all()
            session.close()
            return results
            
        except Exception as e:
            # 获取股票快照数据失败: {str(e)}
            session.close()
            return []
    
    @classmethod
    def get_available_dates(cls) -> List[str]:
        """获取可用的交易日期列表"""
        if not cls._db_manager:
            raise ValueError("数据库管理器未设置，请先调用 set_db_manager()")
        
        try:
            session = cls.get_session()
            dates = session.query(cls.trade_date).distinct().order_by(
                cls.trade_date.desc()
            ).all()
            session.close()
            return [date[0] for date in dates]
        except Exception as e:
            # 获取可用交易日期失败: {str(e)}
            return []
    
    @classmethod
    def get_security_list(cls, trade_date: str) -> List[str]:
        """获取指定日期的股票代码列表"""
        if not cls._db_manager:
            raise ValueError("数据库管理器未设置，请先调用 set_db_manager()")
        
        try:
            session = cls.get_session()
            securities = session.query(cls.security_id).filter(
                cls.trade_date == trade_date
            ).distinct().all()
            session.close()
            return [sec[0] for sec in securities]
        except Exception as e:
            # 获取股票列表失败: {str(e)}
            return []
    
    @classmethod
    def import_from_dataframe(cls, df: pd.DataFrame) -> int:
        """从DataFrame导入数据"""
        if not cls._db_manager:
            raise ValueError("数据库管理器未设置，请先调用 set_db_manager()")

        try:
            session = cls.get_session()
            success_count = 0
            error_count = 0
            
            for index, row in df.iterrows():
                try:
                    # 获取必要的字段值并标准化日期格式
                    trade_date = standardize_date_format(row['trade_date'])
                    market = row['market']
                    security_id = row['security_id']
                    
                    if not trade_date or not security_id:
                        error_count += 1
                        continue
                    
                    # 检查是否已存在相同记录
                    existing = session.query(cls).filter(
                        cls.trade_date == trade_date,
                        cls.market == market,
                        cls.security_id == security_id
                    ).first()
                    
                    if existing:
                        # 更新现有记录的所有字段，使用安全转换
                        value = safe_get_field_value(row, 'pre_close_px')
                        if value is not None:
                            existing.pre_close_px = safe_convert_numeric(value) * 0.0001
                        value = safe_get_field_value(row, 'open_px')
                        if value is not None:
                            existing.open_px = safe_convert_numeric(value) * 0.0001
                        value = safe_get_field_value(row, 'high_px')
                        if value is not None:
                            existing.high_px = safe_convert_numeric(value) * 0.0001
                        value = safe_get_field_value(row, 'low_px')
                        if value is not None:
                            existing.low_px = safe_convert_numeric(value) * 0.0001
                        value = safe_get_field_value(row, 'trade_px')
                        if value is not None:
                            existing.trade_px = safe_convert_numeric(value) * 0.0001
                        value = safe_get_field_value(row, 'num_trades')
                        if value is not None:
                            existing.num_trades = safe_convert_int(value)
                        value = safe_get_field_value(row, 'volume_trade')
                        if value is not None:
                            existing.volume_trade = safe_convert_int(value)
                        value = safe_get_field_value(row, 'amount_trade')
                        if value is not None:
                            existing.amount_trade = safe_convert_numeric(value)
                        value = safe_get_field_value(row, 'pre_volume')
                        if value is not None:
                            existing.pre_volume = safe_convert_int(value)
                        value = safe_get_field_value(row, 'mult_vol')
                        if value is not None:
                            existing.mult_vol = safe_convert_int(value)
                        value = safe_get_field_value(row, 'upper_times')
                        if value is not None:
                            existing.upper_times = safe_convert_int(value)
                        value = safe_get_field_value(row, 'upper_bid_ord_qty')
                        if value is not None:
                            existing.upper_bid_ord_qty = safe_convert_int(value)
                        value = safe_get_field_value(row, 'upper_bid_ord_cnt')
                        if value is not None:
                            existing.upper_bid_ord_cnt = safe_convert_int(value)
                        value = safe_get_field_value(row, 'lower_offer_ord_qty')
                        if value is not None:
                            existing.lower_offer_ord_qty = safe_convert_int(value)
                        value = safe_get_field_value(row, 'lower_offer_ord_cnt')
                        if value is not None:
                            existing.lower_offer_ord_cnt = safe_convert_int(value)
                        value = safe_get_field_value(row, 'upper_offer_ord_qty')
                        if value is not None:
                            existing.upper_offer_ord_qty = safe_convert_int(value)
                        value = safe_get_field_value(row, 'upper_offer_ord_cnt')
                        if value is not None:
                            existing.upper_offer_ord_cnt = safe_convert_int(value)
                        value = safe_get_field_value(row, 'lower_bid_ord_qty')
                        if value is not None:
                            existing.lower_bid_ord_qty = safe_convert_int(value)
                        value = safe_get_field_value(row, 'lower_bid_ord_cnt')
                        if value is not None:
                            existing.lower_bid_ord_cnt = safe_convert_int(value)
                        value = safe_get_field_value(row, 'yes_upper_offer_ord_qty')
                        if value is not None:
                            existing.yes_upper_offer_ord_qty = safe_convert_int(value)
                        value = safe_get_field_value(row, 'yes_upper_offer_ord_cnt')
                        if value is not None:
                            existing.yes_upper_offer_ord_cnt = safe_convert_int(value)
                        value = safe_get_field_value(row, 'upper_px')
                        if value is not None:
                            existing.upper_px = safe_convert_numeric(value) * 0.0001
                        value = safe_get_field_value(row, 'lower_px')
                        if value is not None:
                            existing.lower_px = safe_convert_numeric(value) * 0.0001
                        value = safe_get_field_value(row, 'security_name')
                        if value is not None:
                            existing.security_name = str(value)
                        value = safe_get_field_value(row, 'net_amt')
                        if value is not None:
                            existing.net_amt = safe_convert_numeric(value)
                        value = safe_get_field_value(row, 'big_net_amt')
                        if value is not None:
                            existing.big_net_amt = safe_convert_numeric(value)
                        value = safe_get_field_value(row, 'bias_cnt')
                        if value is not None:
                            existing.bias_cnt = safe_convert_int(value)
                        value = safe_get_field_value(row, 'auction_amt')
                        if value is not None:
                            existing.auction_amt = safe_convert_numeric(value)
                    else:
                        # 创建新记录，使用安全转换
                        snapshot = cls(
                            trade_date=trade_date,
                            market=market,
                            security_id=security_id,
                            pre_close_px=safe_convert_numeric(safe_get_field_value(row, 'pre_close_px', 0)) * 0.0001,
                            open_px=safe_convert_numeric(safe_get_field_value(row, 'open_px', 0)) * 0.0001,
                            high_px=safe_convert_numeric(safe_get_field_value(row, 'high_px', 0)) * 0.0001,
                            low_px=safe_convert_numeric(safe_get_field_value(row, 'low_px', 0)) * 0.0001,
                            trade_px=safe_convert_numeric(safe_get_field_value(row, 'trade_px', 0)) * 0.0001,
                            num_trades=safe_convert_int(safe_get_field_value(row, 'num_trades', 0)),
                            volume_trade=safe_convert_int(safe_get_field_value(row, 'volume_trade', 0)),
                            amount_trade=safe_convert_numeric(safe_get_field_value(row, 'amount_trade', 0)),
                            pre_volume=safe_convert_int(safe_get_field_value(row, 'pre_volume', 0)),
                            mult_vol=safe_convert_int(safe_get_field_value(row, 'mult_vol', 0)),
                            upper_times=safe_convert_int(safe_get_field_value(row, 'upper_times', 0)),
                            upper_bid_ord_qty=safe_convert_int(safe_get_field_value(row, 'upper_bid_ord_qty', 0)),
                            upper_bid_ord_cnt=safe_convert_int(safe_get_field_value(row, 'upper_bid_ord_cnt', 0)),
                            lower_offer_ord_qty=safe_convert_int(safe_get_field_value(row, 'lower_offer_ord_qty', 0)),
                            lower_offer_ord_cnt=safe_convert_int(safe_get_field_value(row, 'lower_offer_ord_cnt', 0)),
                            upper_offer_ord_qty=safe_convert_int(safe_get_field_value(row, 'upper_offer_ord_qty', 0)),
                            upper_offer_ord_cnt=safe_convert_int(safe_get_field_value(row, 'upper_offer_ord_cnt', 0)),
                            lower_bid_ord_qty=safe_convert_int(safe_get_field_value(row, 'lower_bid_ord_qty', 0)),
                            lower_bid_ord_cnt=safe_convert_int(safe_get_field_value(row, 'lower_bid_ord_cnt', 0)),
                            yes_upper_offer_ord_qty=safe_convert_int(safe_get_field_value(row, 'yes_upper_offer_ord_qty', 0)),
                            yes_upper_offer_ord_cnt=safe_convert_int(safe_get_field_value(row, 'yes_upper_offer_ord_cnt', 0)),
                            upper_px=safe_convert_numeric(safe_get_field_value(row, 'upper_px', 0)) * 0.0001,
                            lower_px=safe_convert_numeric(safe_get_field_value(row, 'lower_px', 0)) * 0.0001,
                            security_name=str(safe_get_field_value(row, 'security_name', '') or ''),
                            net_amt=safe_convert_numeric(safe_get_field_value(row, 'net_amt', 0)),
                            big_net_amt=safe_convert_numeric(safe_get_field_value(row, 'big_net_amt', 0)),
                            bias_cnt=safe_convert_int(safe_get_field_value(row, 'bias_cnt', 0)),
                            auction_amt=safe_convert_numeric(safe_get_field_value(row, 'auction_amt', 0))
                        )
                        session.add(snapshot)
                    success_count += 1
                    
                except Exception as e:
                    # 处理第 {index+1} 行数据失败: {str(e)}
                    error_count += 1
                    continue
            
            session.commit()
            session.close()
            # HSnapshot导入完成: {trade_date} 成功 {success_count} 条，失败 {error_count} 条
            return success_count
            
        except Exception as e:
            # 导入股票快照数据失败: {str(e)}
            if 'session' in locals():
                session.rollback()
                session.close()
            return 0
    
    @classmethod
    def to_dataframe(cls, trade_date: Optional[str] = None, execution_date: Optional[str] = None, for_analysis: bool = False) -> pd.DataFrame:
        """转换为DataFrame"""
        if not cls._db_manager:
            raise ValueError("数据库管理器未设置，请先调用 set_db_manager()")
        
        try:
            session = cls.get_session()
            
            query = session.query(cls)
            
            if trade_date:
                query = query.filter(cls.trade_date == trade_date)
            elif execution_date:
                if for_analysis:
                    query = query.filter(cls.trade_date >= execution_date)
                else:
                    query = query.filter(cls.trade_date <= execution_date)
            
            results = query.all()
            
            # 转换为DataFrame
            data_list = []
            for result in results:
                data_list.append(cls._to_dict(result))
            
            session.close()
            
            df = pd.DataFrame(data_list)
            # 从数据库加载了数据
            return df
            
        except Exception as e:
            # 转换为DataFrame失败
            session.close()
            return pd.DataFrame()
    
    @classmethod
    def load_from_db(cls, db_path: str = "stock_data.db", execution_date: Optional[str] = None, for_analysis: bool = False) -> pd.DataFrame:
        """从数据库加载股票快照数据（兼容方法）"""
        # 创建临时数据库管理器
        db_manager = DatabaseManager(stock_data_db_path=db_path)
        cls.set_db_manager(db_manager)
        
        try:
            df = cls.to_dataframe(execution_date=execution_date, for_analysis=for_analysis)
            return df
        finally:
            db_manager.close_all_sessions()
    
    @classmethod
    def preprocess_data(cls, df: pd.DataFrame) -> pd.DataFrame:
        """预处理股票数据"""
        if df.empty:
            return df
        
        try:
            # 确保必要的列存在
            required_columns = ['trade_date', 'security_id', 'trade_px', 'volume_trade']
            missing_columns = [col for col in required_columns if col not in df.columns]
            if missing_columns:
                # 警告: 缺少必要的列
                return df
            
            # 数据类型转换 - 使用copy()避免SettingWithCopyWarning
            df = df.copy()
            if 'trade_date' in df.columns:
                df['trade_date'] = df['trade_date'].astype(str)
            
            if 'security_id' in df.columns:
                df['security_id'] = df['security_id'].astype(str)
            
            # 数值列转换
            numeric_columns = ['trade_px', 'volume_trade', 'amount_trade', 'open_px', 'high_px', 'low_px','pre_close_px']
            for col in numeric_columns:
                if col in df.columns:
                    df[col] = pd.to_numeric(df[col], errors='coerce')

            # 为策略与前端使用创建通用列别名
            alias_mapping = {
                'open': 'open_px',
                'high': 'high_px',
                'low': 'low_px',
                'close': 'trade_px',
                'volume': 'volume_trade',
                'pre_close': 'pre_close_px'
            }
            for alias, source in alias_mapping.items():
                if alias not in df.columns and source in df.columns:
                    df[alias] = pd.to_numeric(df[source], errors='coerce')
            
            # 去除重复数据
            original_count = len(df)
            df = df.drop_duplicates(subset=['trade_date', 'security_id'], keep='last')
            deduplicated_count = len(df)
            
            if original_count != deduplicated_count:
                pass  # 数据去重完成
            
            # 按日期和股票代码排序
            df = df.sort_values(['trade_date', 'security_id'])
            
            # 数据预处理完成
            return df
            
        except Exception as e:
            # 数据预处理失败
            return df
